#!/bin/bash
# ADKR Operations Manager Script
# Provides comprehensive operations management for ADKR deployment

set -e

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
DEFAULT_CONFIG_DIR="./keys"
DEFAULT_BINARY_PATH="./adkr_node"
DEFAULT_HOSTS_FILE="hosts.template"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Help function
show_help() {
    cat << EOF
ADKR Operations Manager

Usage: $0 <command> [options]

Commands:
    setup           Set up ADKR environment
    generate-keys   Generate BLS keys and configuration files
    validate        Validate configuration files
    deploy          Deploy ADKR nodes
    start           Start ADKR nodes
    stop            Stop ADKR nodes
    status          Check status of ADKR nodes
    logs            Show logs from ADKR nodes
    clean           Clean up ADKR deployment
    benchmark       Run performance benchmark
    test            Run equivalence tests

Options:
    -n, --nodes     Number of nodes (default: 4)
    -t, --threshold Threshold for BLS keys (default: 3)
    -c, --config    Configuration directory (default: ./keys)
    -b, --binary    Path to adkr_node binary (default: ./adkr_node)
    -h, --hosts     Hosts configuration file (default: hosts.template)
    -f, --fast      Use fast mode for benchmarks
    -v, --verbose   Verbose output

Examples:
    $0 setup -n 4 -t 3
    $0 generate-keys -n 4 -t 3 -c ./mykeys
    $0 validate -c ./keys
    $0 deploy -c ./keys
    $0 start -c ./keys
    $0 benchmark -f -n 4
    $0 test

EOF
}

# Parse command line arguments
parse_args() {
    COMMAND=""
    NODES=4
    THRESHOLD=3
    CONFIG_DIR="$DEFAULT_CONFIG_DIR"
    BINARY_PATH="$DEFAULT_BINARY_PATH"
    HOSTS_FILE="$DEFAULT_HOSTS_FILE"
    FAST_MODE=false
    VERBOSE=false
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            setup|generate-keys|validate|deploy|start|stop|status|logs|clean|benchmark|test)
                COMMAND="$1"
                shift
                ;;
            -n|--nodes)
                NODES="$2"
                shift 2
                ;;
            -t|--threshold)
                THRESHOLD="$2"
                shift 2
                ;;
            -c|--config)
                CONFIG_DIR="$2"
                shift 2
                ;;
            -b|--binary)
                BINARY_PATH="$2"
                shift 2
                ;;
            -h|--hosts)
                HOSTS_FILE="$2"
                shift 2
                ;;
            -f|--fast)
                FAST_MODE=true
                shift
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
    
    if [[ -z "$COMMAND" ]]; then
        log_error "No command specified"
        show_help
        exit 1
    fi
}

# Setup ADKR environment
setup_environment() {
    log_info "Setting up ADKR environment..."
    
    # Check if Go is installed
    if ! command -v go &> /dev/null; then
        log_error "Go is not installed. Please install Go 1.19 or later."
        exit 1
    fi
    
    # Check Go version
    GO_VERSION=$(go version | awk '{print $3}' | sed 's/go//')
    log_info "Go version: $GO_VERSION"
    
    # Build ADKR binaries
    log_info "Building ADKR binaries..."
    cd "$PROJECT_ROOT"
    
    if ! go build -o adkr_node ./cmd/adkr_node; then
        log_error "Failed to build adkr_node"
        exit 1
    fi
    
    if ! go build -o adkr_bench ./cmd/adkr_bench; then
        log_error "Failed to build adkr_bench"
        exit 1
    fi
    
    if ! go build -o bls_gen ./cmd/bls_gen; then
        log_error "Failed to build bls_gen"
        exit 1
    fi
    
    if ! go build -o config_validator ./cmd/config_validator; then
        log_error "Failed to build config_validator"
        exit 1
    fi
    
    log_success "ADKR environment setup completed"
}

# Generate BLS keys and configuration
generate_keys() {
    log_info "Generating BLS keys and configuration files..."
    log_info "Nodes: $NODES, Threshold: $THRESHOLD, Config Dir: $CONFIG_DIR"
    
    cd "$PROJECT_ROOT"
    
    # Generate keys
    if ! ./bls_gen -n "$NODES" -t "$THRESHOLD" -output "$CONFIG_DIR"; then
        log_error "Failed to generate BLS keys"
        exit 1
    fi
    
    # Validate generated keys
    if ! ./config_validator -config "$CONFIG_DIR"; then
        log_error "Key validation failed"
        exit 1
    fi
    
    log_success "BLS keys and configuration generated successfully"
    log_info "Configuration files created in: $CONFIG_DIR"
}

# Validate configuration
validate_config() {
    log_info "Validating ADKR configuration..."
    log_info "Config Dir: $CONFIG_DIR"
    
    cd "$PROJECT_ROOT"
    
    if ! ./config_validator -config "$CONFIG_DIR"; then
        log_error "Configuration validation failed"
        exit 1
    fi
    
    log_success "Configuration validation passed"
}

# Deploy ADKR nodes
deploy_nodes() {
    log_info "Deploying ADKR nodes..."
    log_info "Config Dir: $CONFIG_DIR, Binary: $BINARY_PATH"
    
    cd "$PROJECT_ROOT"
    
    # Check if configuration exists
    if [[ ! -d "$CONFIG_DIR" ]]; then
        log_error "Configuration directory not found: $CONFIG_DIR"
        log_info "Run 'generate-keys' command first"
        exit 1
    fi
    
    # Check if binary exists
    if [[ ! -f "$BINARY_PATH" ]]; then
        log_error "Binary not found: $BINARY_PATH"
        log_info "Run 'setup' command first"
        exit 1
    fi
    
    # Create deployment script
    DEPLOY_SCRIPT="$CONFIG_DIR/deploy.sh"
    if [[ -f "$DEPLOY_SCRIPT" ]]; then
        log_info "Using existing deployment script: $DEPLOY_SCRIPT"
        chmod +x "$DEPLOY_SCRIPT"
    else
        log_info "Creating deployment script..."
        cat > "$DEPLOY_SCRIPT" << EOF
#!/bin/bash
set -e

BINARY_PATH="$BINARY_PATH"
HOSTS_FILE="$HOSTS_FILE"
CONFIG_DIR="$CONFIG_DIR"

echo "Starting ADKR deployment..."

EOF
        
        for i in $(seq 0 $((NODES-1))); do
            cat >> "$DEPLOY_SCRIPT" << EOF
echo "Starting node $i..."
\$BINARY_PATH -pid $i -hosts "\$CONFIG_DIR/\$HOSTS_FILE" -key "\$CONFIG_DIR/bls_keys_node_$i.json" &
sleep 1

EOF
        done
        
        cat >> "$DEPLOY_SCRIPT" << EOF
echo "All nodes started. Use 'pkill adkr_node' to stop."
EOF
        
        chmod +x "$DEPLOY_SCRIPT"
    fi
    
    log_success "Deployment script created: $DEPLOY_SCRIPT"
    log_info "Run '$DEPLOY_SCRIPT' to start nodes"
}

# Start ADKR nodes
start_nodes() {
    log_info "Starting ADKR nodes..."
    
    cd "$PROJECT_ROOT"
    
    DEPLOY_SCRIPT="$CONFIG_DIR/deploy.sh"
    if [[ ! -f "$DEPLOY_SCRIPT" ]]; then
        log_error "Deployment script not found: $DEPLOY_SCRIPT"
        log_info "Run 'deploy' command first"
        exit 1
    fi
    
    # Check if nodes are already running
    if pgrep -f "adkr_node" > /dev/null; then
        log_warning "ADKR nodes appear to be already running"
        log_info "Use 'stop' command to stop existing nodes first"
        exit 1
    fi
    
    # Start nodes
    "$DEPLOY_SCRIPT"
    
    # Wait a moment and check status
    sleep 2
    if pgrep -f "adkr_node" > /dev/null; then
        log_success "ADKR nodes started successfully"
    else
        log_error "Failed to start ADKR nodes"
        exit 1
    fi
}

# Stop ADKR nodes
stop_nodes() {
    log_info "Stopping ADKR nodes..."
    
    if ! pgrep -f "adkr_node" > /dev/null; then
        log_warning "No ADKR nodes are running"
        return 0
    fi
    
    # Stop all adkr_node processes
    pkill -f "adkr_node" || true
    
    # Wait for processes to stop
    sleep 2
    
    if ! pgrep -f "adkr_node" > /dev/null; then
        log_success "ADKR nodes stopped successfully"
    else
        log_warning "Some ADKR nodes may still be running"
    fi
}

# Check status of ADKR nodes
check_status() {
    log_info "Checking ADKR node status..."
    
    RUNNING_NODES=$(pgrep -f "adkr_node" | wc -l)
    
    if [[ $RUNNING_NODES -eq 0 ]]; then
        log_info "No ADKR nodes are running"
    else
        log_info "Found $RUNNING_NODES running ADKR node(s)"
        
        if [[ "$VERBOSE" == "true" ]]; then
            log_info "Running processes:"
            pgrep -f "adkr_node" | while read pid; do
                ps -p "$pid" -o pid,ppid,cmd --no-headers
            done
        fi
    fi
}

# Show logs from ADKR nodes
show_logs() {
    log_info "Showing ADKR node logs..."
    
    if ! pgrep -f "adkr_node" > /dev/null; then
        log_warning "No ADKR nodes are running"
        return 0
    fi
    
    # Show recent logs (this would need to be implemented based on logging setup)
    log_info "Log output from running nodes:"
    # For now, just show process information
    pgrep -f "adkr_node" | while read pid; do
        echo "Node PID: $pid"
    done
}

# Clean up ADKR deployment
clean_deployment() {
    log_info "Cleaning up ADKR deployment..."
    
    # Stop nodes
    stop_nodes
    
    # Remove log files (if any)
    find "$PROJECT_ROOT" -name "*.log" -delete 2>/dev/null || true
    
    # Remove temporary files
    find "$PROJECT_ROOT" -name "*.tmp" -delete 2>/dev/null || true
    
    log_success "Cleanup completed"
}

# Run performance benchmark
run_benchmark() {
    log_info "Running ADKR performance benchmark..."
    
    cd "$PROJECT_ROOT"
    
    BENCH_ARGS="-n $NODES"
    if [[ "$FAST_MODE" == "true" ]]; then
        BENCH_ARGS="$BENCH_ARGS -fast"
    fi
    
    if [[ "$VERBOSE" == "true" ]]; then
        BENCH_ARGS="$BENCH_ARGS -v"
    fi
    
    if ! ./adkr_bench $BENCH_ARGS; then
        log_error "Benchmark failed"
        exit 1
    fi
    
    log_success "Benchmark completed"
}

# Run equivalence tests
run_tests() {
    log_info "Running ADKR equivalence tests..."
    
    cd "$PROJECT_ROOT"
    
    # Run simplified parity tests
    if ! go build -o parity_simple ./cmd/parity_simple; then
        log_error "Failed to build parity_simple"
        exit 1
    fi
    
    if ! ./parity_simple -output parity_test_results.json; then
        log_error "Parity tests failed"
        exit 1
    fi
    
    # Run analysis
    if command -v python3 &> /dev/null; then
        if [[ -f "parity_analysis.py" ]]; then
            python3 parity_analysis.py parity_test_results.json
        fi
    fi
    
    log_success "Equivalence tests completed"
}

# Main function
main() {
    parse_args "$@"
    
    log_info "ADKR Operations Manager"
    log_info "Command: $COMMAND"
    
    case "$COMMAND" in
        setup)
            setup_environment
            ;;
        generate-keys)
            generate_keys
            ;;
        validate)
            validate_config
            ;;
        deploy)
            deploy_nodes
            ;;
        start)
            start_nodes
            ;;
        stop)
            stop_nodes
            ;;
        status)
            check_status
            ;;
        logs)
            show_logs
            ;;
        clean)
            clean_deployment
            ;;
        benchmark)
            run_benchmark
            ;;
        test)
            run_tests
            ;;
        *)
            log_error "Unknown command: $COMMAND"
            show_help
            exit 1
            ;;
    esac
}

# Run main function
main "$@"



